"""
来自论文[1] 许利恒.基于椭圆特征的陨石坑检测与识别方法研究[D].北京:北京航空航天大学,2022

按照论文的方法复现，关于从四点金字塔拓展至整幅陨石坑的算法，这里使用了姿态解算的结果作为重投影平差的依据，不再使用导航坑对的不变量计算，增加了对陨石坑的筛选算法
"""

import numpy as np
from typing import Iterator
from ...base import Matching
from utils.ellipse import ellipse_args, center_ellipse
from itertools import combinations
from scipy.optimize import linear_sum_assignment
from utils.pose import pose_calculate
import cv2


class BinadFastPyramid(Matching):
    def __init__(self, catalog_dir, **kwargs):
        super().__init__(catalog_dir, **kwargs)

    def binad_descriptor(
        self, C1, C2, *args, **kwargs
    ) -> tuple[np.ndarray, np.ndarray]:
        """从陨石坑的参数中计算描述子，C1和C2可以是含有不确定度的椭圆参数"""
        raise NotImplementedError

    def extend_identify(self, params, ijkl, IJKL, *args, **kwargs) -> tuple[bool, list]:
        raise NotImplementedError

    def identify(
        self, params, uncertainty, factor, *args, max_try=5, **kwargs
    ) -> Iterator:
        """
        输入应该越接近于真实的检测输出越好。因此输入应当是一堆椭圆的拟合参数，输出应当是匹配确定的结果，即椭圆ID的列表。
        Arguments:
            params (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1)
        Returns:
            list : matched ellipse ID
        """
        # 将陨石坑按直径排序
        params = np.array(params)
        elli = ellipse_args(params)
        diameters = elli[2] + elli[3]
        idx = np.argsort(-diameters)
        # params = params[idx]
        # 似乎也可以不用构建全部二元组
        # 从直径最大、且二元组最多的陨石坑对开始构建金字塔
        not_used: list = idx.tolist()
        # 试错次数可以作为超参数设置
        while len(not_used) > 3:
            i, j = not_used[:2]
            fail_num = 0
            for k, l in combinations(not_used[2:], 2):
                # 批量算出金字塔的三元组
                I12, th = self.binad_descriptor(
                    params[[i, j, i, i, j, k]],
                    params[[j, k, k, l, l, l]],
                    uncertainty[[i, j, i, i, j, k]],
                    uncertainty[[j, k, k, l, l, l]],
                    factor,
                )
                AB = self.binad_search(I12, th)
                ABC = self.build_triads(*AB[:3])
                for IJKL in self.build_pyramids(*AB[3:], ABC):
                    # 完成一次全图扩展
                    success, result = self.extend_identify(
                        params,
                        uncertainty,
                        (i, j, k, l),
                        IJKL,
                        *args,
                        **kwargs,
                    )
                    if success:
                        return result
                fail_num += 1
                if fail_num == max_try:
                    break
            not_used.pop(0)
        return None

    def binad_search(self, binad, th=None):
        """
        从库内搜索阈值范围内的全部二元组
        Arguments:
            binad (np.ndarray) : 二元组的特征向量
            th (float) : 搜索阈值，形状应当与描述子相同
        """
        assert th.shape == binad.shape
        desc = self.descriptor[:, None].repeat(binad.shape[1], axis=1)
        # 正向查找
        ind_1 = (np.abs(desc - binad[..., None]) < th[..., None]).all(axis=0)
        # 反向查找
        ind_2 = (np.abs(desc - binad[::-1, :, None]) < th[::-1, :, None]).all(axis=0)
        # 所有特征维度压缩至同一个维度，使用&操作符
        desc = []
        for id_1, id_2 in zip(ind_1, ind_2):
            bin = np.concatenate((self.idx[:, id_1], self.idx[:, id_2][::-1]), axis=1).T
            desc.append(bin)
        return desc

    def build_triads(self, IJ, JK, IK):
        # 找出最短的二元组
        min_len = np.array([len(IJ), len(JK), len(IK)]).argmin()
        IJK = []
        if min_len == 0:
            for i, j in IJ:
                k = np.intersect1d(JK[JK[:, 0] == j, 1], IK[IK[:, 0] == i, 1])
                if k.size > 0:
                    IJK.append(
                        np.vstack((i * np.ones_like(k), j * np.ones_like(k), k)).T
                    )
        elif min_len == 1:
            for j, k in JK:
                i = np.intersect1d(IJ[IJ[:, 1] == j, 0], IK[IK[:, 1] == k, 0])
                if i.size > 0:
                    IJK.append(
                        np.vstack((i, j * np.ones_like(i), k * np.ones_like(i))).T
                    )
        else:
            for i, k in IK:
                j = np.intersect1d(JK[JK[:, 1] == k, 0], IJ[IJ[:, 0] == i, 1])
                if j.size > 0:
                    IJK.append(
                        np.vstack((i * np.ones_like(j), j, k * np.ones_like(j))).T
                    )
        if len(IJK) == 0:
            return []
        return np.concatenate(IJK, axis=0)

    def build_pyramids(self, IL, JL, KL, IJK):
        """
        从库内搜索阈值范围内的全部四元组
        Arguments:
            pyramid (np.ndarray) : 四元组的特征向量
            th (float) : 搜索阈值，形状应当与描述子相同
        """
        # 查找ijls中的ij
        for I, J, K in IJK:
            ind_I = IL[IL[:, 0] == I, 1]
            ind_J = JL[JL[:, 0] == J, 1]
            ind_K = KL[KL[:, 0] == K, 1]
            Ls = np.intersect1d(np.intersect1d(ind_I, ind_J), ind_K)
            if Ls.size > 0:
                for L in Ls:
                    yield I, J, K, L

    def extend_identify(
        self,
        params,
        uncertainty,
        ijkl,
        IJKL,
        K,
        th: float,
        confidence: float,
        dist_th=2,
        **kwargs,
    ):
        """
        TODO 用增量式重投影方法完成对全幅陨石坑的匹配，具体的细节如下：
        1. 利用现有四点ijkl估计出一个摄像机矩阵P（TODO 有可能估计失败）
        2. 根据当前的四点ijkl，设置某个超参数r，搜索四点周围r个点构成候选重投影点，利用P重投影至图像平面，记为Qr
        3. 利用匈牙利算法，对当前的Qr+4个重投影点与当前图像上的全部陨石坑中心点（可能多于Qr+4）作匹配（TODO 算法输出差过大的点应当舍去）。
        4. 对于每个匹配上的陨石坑中心点，计算其不变量以排除潜在的噪声点。
        4. 匹配的结果作为新的ijkl，重新估计P，重复步骤2-3，直到全部匹配上或者达到最大迭代次数
        """
        # 返回第一符合要求的值
        # 将金字塔识别结果拓展至整个图像
        if len(ijkl) < 4:
            return False, None
        cnts = np.array(center_ellipse(params)).T
        ijkl, IJKL = self.increment_hungary(params, uncertainty, ijkl, IJKL, K, dist_th)
        if len(ijkl) < 4:
            return False, None
            # 重投影验证
        points_3d = self.catalog[IJKL, :]
        obse_2d = cnts[ijkl, :]
        R, T = pose_calculate(K, points_3d, obse_2d)
        if R is None:
            return False, None
        uv = cv2.projectPoints(points_3d, R, T, K, np.zeros(4))[0].squeeze()
        ind = np.linalg.norm(uv - obse_2d, axis=1) < th
        if ind.mean() >= confidence:
            return True, (ijkl[ind], IJKL[ind])
        else:
            return False, None

    def increment_hungary(
        self,
        params,
        uncertainty,
        ijkl: tuple,
        IJKL: tuple,
        K,
        filter_invt=True,
        dist_th=2.0,
        confidence=0.8,
    ):
        """
        用增量式重投影方法完成对全幅陨石坑的匹配，具体的细节如下：
        1. 利用现有四点ijkl估计出一个摄像机矩阵P（TODO 有可能估计失败）
        2. 根据当前的四点ijkl，设置某个超参数r，搜索四点周围r个点构成候选重投影点，利用P重投影至图像平面，记为Qr
        3. 利用匈牙利算法，对当前的Qr+4个重投影点与当前图像上的全部陨石坑中心点（可能多于Qr+4）作匹配；
        4. 对于每个匹配上的陨石坑中心点，计算其不变量以排除潜在的噪声点。
        4. 匹配的结果作为新的ijkl，重新估计P，重复步骤2-3，直到全部匹配上或者达到最大迭代次数
        Arguments:
            ijkl (tuple) : 当前图像每个被匹配的陨石坑序号
            IJKL (tuple) : 当前图像每个被匹配的陨石坑在目录中的序号
            K (np.ndarray) : 相机内参矩阵
            filter_invt (bool) : 是否使用不变量验证
            dist_th (float) : 仅在filter_invt=True时有效匹配距离阈值
            factor (float) : 仅在filter_invt=True时有效，不变量验证的阈值
            confidence (float) : 仅在filter_invt=True时有效，用于筛选错误匹配的置信概率，推荐取值不超过0.9，低噪声情况下可适当取低如0.7甚至0.65均可取得相当好的效果
        """
        # 估计P
        points_3d = self.catalog[IJKL, :]
        cnts = np.array(center_ellipse(params)).T
        obse_2d = cnts[ijkl, :]
        R, T = pose_calculate(K, points_3d, obse_2d)
        if R is None:
            return [], []
        # 重投影
        uv = cv2.projectPoints(self.catalog, R, T, K, np.zeros(4))[0].squeeze()
        # 把当前可能位于图像平面内的全部点都加入至候选列表中
        ind = (
            (uv[:, 0] < K[0, 2] * 2)
            & (uv[:, 0] >= 0)
            & (uv[:, 1] < K[1, 2] * 2)
            & (uv[:, 1] >= 0)
        )
        # 得到全部候选中心点列表
        cata_3d = uv[ind]
        cata_id = np.arange(self.catalog.shape[0])[ind]
        # 获取当前全部图像中心点
        obse_id = np.arange(len(cnts))
        # 匈牙利算法
        dist_matrix = np.linalg.norm(cata_3d[:, np.newaxis] - cnts, axis=2)
        # 使用匈牙利算法计算最优分配
        row_ind, col_ind = linear_sum_assignment(dist_matrix)
        # 最优分配
        cata_id = cata_id[row_ind]
        obse_id = obse_id[col_ind]
        # 计算匹配误差，本误差也将作为筛选真实陨石坑和虚假陨石坑的依据
        error = dist_matrix[row_ind, col_ind]
        # 计算不变量验证
        ## 具体的步骤应当包括：
        ## 1. 根据匈牙利算法匹配出的距离远近和设计的超参数dist_th，选出符合匹配距离范围内的候选陨石坑
        ## 2. 全部候选陨石坑两两作对，算出不变量以及配对
        ## 3. 根据配对关系，给出导航坑表内的不变量对
        # 挑选匹配距离范围内的陨石坑
        ind = error < dist_th
        cata_id = cata_id[ind]
        obse_id = obse_id[ind]
        if filter_invt:
            # 新增一个哈希列表暂存每对陨石坑错误匹配的次数
            err_num = np.zeros_like(obse_id, dtype=np.int16)
            # 计算容错数量阈值，该值为 n - 0.5 - sqrt((th*n^2-th*n+0.25)，计算过程见于汇报2024年9月5日
            num_th = obse_id.shape[0]
            num_th = (
                num_th
                - 0.5
                - np.sqrt(confidence * num_th * num_th - confidence * num_th + 0.25)
            )
            if obse_id.shape[0] < 4:
                return [], []
            # 批量完成不变量验证
            ij = np.array(list(combinations(range(obse_id.shape[0]), 2)))
            i12, th = self.binad_descriptor(
                params[obse_id[ij[:, 0]]], params[obse_id[ij[:, 1]]]
            )
            idx = np.argwhere(
                (
                    self.idx[0, None].repeat(ij.shape[1], axis=0) - cata_id[ij[0], None]
                    == 0
                )
                & (
                    self.idx[1, None].repeat(ij.shape[1], axis=0) - cata_id[ij[1], None]
                    == 0
                )
            )
            i12 = i12[:, idx[:, 0]]
            I12 = self.descriptor[:, idx[:, 1]]
            err_num[ij[0]] = (
                err_num[ij[0]]
                + (
                    (np.abs(i12 - I12) > th[:, idx[:, 0]]).sum(axis=0)
                    > (1 - confidence**2) * th.shape[0]
                ).sum()
            )
            err_num[ij[1]] = (
                err_num[ij[1]]
                + (
                    (np.abs(i12 - I12) > th[:, idx[:, 0]]).sum(axis=0)
                    > (1 - confidence**2) * th.shape[0]
                ).sum()
            )
            # 删除错误匹配
            ind = err_num < num_th
            cata_id = cata_id[ind]
            obse_id = obse_id[ind]
        return obse_id, cata_id
